
import numpy as np
import cv2

cap=cv2.VideoCapture(r'../video/青岛胶州湾跨海大桥.FLV')


while not cap.isOpened():
    print('open video error!')
    exit()

#Shi-Tomasi角点检测的参数
feature_params=dict(maxCorners=100,
                    qualityLevel=0.3,
                    minDistance=7,
                    blockSize=7)

#Lucas-Kanade光流算法的参数
lk_params=dict(winSize=(15,15), #窗口大小
               maxLevel=2,
               criteria=(cv2.TERM_CRITERIA_EPS| cv2.TERM_CRITERIA_COUNT,10,0.03))
#创建一些随机颜色
color=np.random.randint(0,255,(100,3))

#取第一帧，检测其中的角点
ret,old_frame=cap.read()
old_gray=cv2.cvtColor(old_frame,cv2.COLOR_BGR2GRAY)
##p0的形式：[  [[x1,y1]], [[x2,y2]], ... ,[[xn,yn]]   ]
p0 =cv2.goodFeaturesToTrack(old_gray, mask=None, **feature_params)

#创建一个掩膜图像做绘制用
mask=np.zeros_like(old_frame)   #返回0阵列，其尺度跟old_frame一样

while(1):
    ret,frame=cap.read()
    frame_gray=cv2.cvtColor(frame,cv2.COLOR_BGR2GRAY)

    ##计算光流，
    ##p1和点数量和p0一样
    ##st是和p1,p0数量一样的向量，1表示p0对应的p1找到了
    p1,st,err=cv2.calcOpticalFlowPyrLK(old_gray,frame_gray,p0,None,**lk_params)
    #选择好的点
    good_new=p1[st==1]  #取好的匹配  得到的形式：[ [x1,y1],[x2,y2], ... ,[xn,yn]  ]
    good_old=p0[st==1]

    #绘制光流轨迹
    for i,(new,old) in enumerate(zip(good_new,good_old)):
        a,b=new.ravel() #阵列平坦化
        c,d=old.ravel()
        #绘制光流方向线段
        mask=cv2.line(mask,(a,b),(c,d),color[i].tolist(),2)
        frame=cv2.circle(frame,(a,b),5,color[i].tolist(),-1)
    img=cv2.add(frame,mask)

    cv2.imshow('frame',img)
    k=cv2.waitKey(30) & 0xff
    if k==27:
        break

    ##更新前一帧，以及前一帧的点
    old_gray=frame_gray.copy()
    p0=good_new.reshape(-1,1,2)

cv2.destroyAllWindows()
cap.release()